import pandas as pd
from datetime import datetime


def get_quarter(date_obj):
    """
    根据日期对象（pandas.Timestamp）返回季度。
    :param date_obj: 随访日期（pandas.Timestamp 类型）
    :return: 季度（1到4）
    """
    if isinstance(date_obj, str):
        date_obj = datetime.strptime(date_obj, "%Y-%m-%d")

    month = date_obj.month
    if 1 <= month <= 3:
        return 1
    elif 4 <= month <= 6:
        return 2
    elif 7 <= month <= 9:
        return 3
    else:
        return 4


def error3(df, styles_df, column1, column2, column3, column4, column5, column6):
    """
    比较同一身份证号的随访记录，若前后季度相同，则标绿。

    :param df: 原始数据 DataFrame
    :param styles_df: 样式 DataFrame，用于设置 Excel 格式
    :param column1: 身份证号
    :param column2: 随访日期
    :param column3: 收缩压（mmHg）
    :param column4: 舒张压（mmHg）
    :param column5: 心率（次/分）
    :param column6: 体重（kg）
    """
    try:
        # 遍历每一行，处理同一身份证号的随访记录
        def highlight_if_duplicates(row):
            idcard = row[column1]
            visit_date = row[column2]
            systolic = row[column3]
            diastolic = row[column4]
            heart_rate = row[column5]
            weight = row[column6]

            # 如果 visit_date 是 Timestamp 类型，将其转换为日期对象
            if isinstance(visit_date, pd.Timestamp):
                visit_date = visit_date.strftime("%Y-%m-%d")  # 转换为字符串

            # 获取当前行的季度
            current_quarter = get_quarter(visit_date)

            # 找出相同身份证号的所有记录
            same_id_df = df[df[column1] == idcard]

            # 获取当前身份证号的季度随访数据
            for _, other_row in same_id_df.iterrows():
                # 比较不同季度的记录
                if other_row[column2] != visit_date:  # 不比较同一行
                    other_visit_date = other_row[column2]

                    # 如果其他行的日期是 Timestamp 类型，也转换为字符串
                    if isinstance(other_visit_date, pd.Timestamp):
                        other_visit_date = other_visit_date.strftime("%Y-%m-%d")

                    other_quarter = get_quarter(other_visit_date)

                    # 判断是否在相邻季度之间有相同数据
                    if abs(current_quarter - other_quarter) == 1:  # 相邻季度
                        # 检查收缩压、舒张压、心率、体重是否相同
                        if (systolic == other_row[column3] and
                                diastolic == other_row[column4] and
                                heart_rate == other_row[column5] and
                                weight == other_row[column6]):
                            # 标记绿色
                            styles_df.loc[row.name, [column3, column4, column5, column6]] = 'background-color: green'
                            styles_df.loc[
                                other_row.name, [column3, column4, column5, column6]] = 'background-color: green'

            return [''] * len(row)

        # 应用样式
        df.apply(highlight_if_duplicates, axis=1)
        print("[随访]已完成处理。")

    except Exception as e:
        print(f"随访处理时发生错误：{e}")
